package com.yeskery.nut.extend.elasticsearch;

import com.yeskery.nut.core.ResponseCode;
import com.yeskery.nut.util.StringUtils;

import java.util.Iterator;
import java.util.Spliterator;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Consumer;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

/**
 * 分页ES集合对象分裂器
 * @author YESKERY
 * 2024/8/29
 */
public class PageListSpliterator<T> implements Spliterator<T> {

    /** ES操作模板 */
    private final ElasticSearchTemplate elasticSearchTemplate;

    /** 索引名称 */
    private final String indexName;

    /** 每页的条数 */
    private final int size;

    /** 总条数 */
    private final long total;

    /** 总页数 */
    private final long pages;

    /** 当前页数 */
    private final AtomicLong currentPage;

    /** 结果类型 */
    private final Class<T> clazz;

    /** 查询语句 */
    private final String query;

    /** 排序语句 */
    private final String sort;

    /** 结果迭代器 */
    private Iterator<T> iterator;

    /**
     * 构建分页ES集合对象分裂器
     * @param elasticSearchTemplate ES操作模板
     * @param indexName 索引名称
     * @param query 查询语句DSL，{"match_all": {}}
     * @param sort 排序语句
     * @param size 每页的条数
     * @param clazz 结果类型
     */
    public PageListSpliterator(ElasticSearchTemplate elasticSearchTemplate, String indexName, String query, String sort, int size, Class<T> clazz) {
        this.elasticSearchTemplate = elasticSearchTemplate;
        this.indexName = indexName;
        if (StringUtils.isEmpty(query)) {
            throw new ElasticSearchException(ResponseCode.BAD_REQUEST.getCode(), "Query Must Not Be Empty.");
        }
        this.query = query;
        this.sort = sort;
        this.size = size;
        this.total = elasticSearchTemplate.countByQuery(indexName, query);
        this.pages = total % size == 0 ? total / size : total / size + 1;
        this.currentPage = new AtomicLong(1);
        this.clazz = clazz;
        this.iterator = getIterator();
    }

    @Override
    public boolean tryAdvance(Consumer<? super T> action) {
        if (total <= 0) {
            return false;
        }
        if (iterator.hasNext()) {
            action.accept(iterator.next());
            return true;
        }
        if (currentPage.longValue() < pages) {
            currentPage.incrementAndGet();

            iterator = getIterator();
            if (iterator.hasNext()) {
                action.accept(iterator.next());
                return true;
            }
        }
        return false;
    }

    @Override
    public Spliterator<T> trySplit() {
        return null;
    }

    @Override
    public long estimateSize() {
        return Long.MAX_VALUE;
    }

    @Override
    public int characteristics() {
        return Spliterator.ORDERED;
    }

    /**
     * 生成Stream流
     * @return Stream流
     */
    public Stream<T> stream() {
        return StreamSupport.stream(this, false);
    }

    /**
     * 获取结果集对象
     * @return 结果集对象
     */
    private Iterator<T> getIterator() {
        return elasticSearchTemplate.searchList(indexName, query, (currentPage.longValue() - 1) * size, size, sort, clazz).iterator();
    }
}
